import os
import random
from PIL import Image
import cv2
from matplotlib import pyplot as plt

val_path = "/home/pxxy2417/CelebA/Img/data/val"

path_list = random.choices(os.listdir(val_path + "/images"), k=5)

xy_dict = {}
xy_dict_original = {}
for path in path_list:
    img = Image.open(f"{val_path}/images/{path}")
    img_w, img_h = img.size
    with open(f"{val_path}/labels/{os.path.splitext(path)[0]}.txt") as f:
        line = f.readline().split()
        x2mx1 = float(line[3]) * img_w
        y2my1 = float(line[4]) * img_h
        x1px2 = float(line[1]) * img_w * 2.0
        y1py2 = float(line[2]) * img_h * 2.0
        x2 = (x1px2 + x2mx1) / 2.0
        x1 = (x1px2 - x2mx1) / 2.0
        y2 = (y1py2 + y2my1) / 2.0
        y1 = (y1py2 - y2my1) / 2.0
    xy_dict[path] = [(int(x1), int(y1)), (int(x2), int(y2))]
    xy_dict_original[path] = [(int(x1), int(y1)), (int(x2), int(y2 / 0.88))]
print(xy_dict)
plt.figure(figsize=(17, 7))
for index, (name, xy) in enumerate(xy_dict.items()):
    image = cv2.imread(os.path.join(val_path, "images", name))
    cv2.rectangle(image, xy[0], xy[1], (255, 255, 255), 2)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    plt.subplot(2, 5, index + 1)
    plt.imshow(image)
    plt.axis('off')
    plt.title(name)

    image = cv2.imread(os.path.join(val_path, "images", name))
    cv2.rectangle(image, xy_dict_original[name][0], xy_dict_original[name][1], (255, 255, 255), 2)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    plt.subplot(2, 5, index + 1 + 5)
    plt.imshow(image)
    plt.axis('off')
    plt.title(name+" original")
plt.show()


